import os
import re
import json
from collections import defaultdict

def extract_cli_node_from_code(code):
    # 去掉注释
    code = re.sub(r'//.*', '', code)
    code = re.sub(r'/\*.*?\*/', '', code, flags=re.DOTALL)

    # 提取 DEFUN 宏
    defun_pattern = re.compile(
        r'DEFUN\s*\(\s*([a-zA-Z0-9_]+)\s*,\s*([a-zA-Z0-9_]+)\s*,\s*"([^"]+)"',
        re.DOTALL
    )
    defun_cmd_map = {}  # handler_cmd -> (cli, handler)
    for m in defun_pattern.finditer(code):
        handler = m.group(1)
        handler_cmd = m.group(2)
        cli = m.group(3)
        defun_cmd_map[handler_cmd] = (cli, handler)

    # 提取 install_element 调用
    install_pattern = re.compile(
        r'install_element\s*\(\s*([A-Z_]+)\s*,\s*&([a-zA-Z0-9_]+)\s*\)',
        re.DOTALL
    )
    node_cmd_map = []  # (NODE, handler_cmd)
    for m in install_pattern.finditer(code):
        node = m.group(1)
        handler_cmd = m.group(2)
        node_cmd_map.append((node, handler_cmd))

    # 关联映射
    cli_node_map = defaultdict(set)  # cli -> set of NODEs
    for node, handler_cmd in node_cmd_map:
        if handler_cmd in defun_cmd_map:
            cli, handler = defun_cmd_map[handler_cmd]
            cli_node_map[cli].add(node)
    return cli_node_map

def scan_folder(folder):
    all_cli_node_map = defaultdict(set)
    for root, dirs, files in os.walk(folder):
        for fname in files:
            if fname.endswith('.c'):
                fpath = os.path.join(root, fname)
                with open(fpath, 'r', encoding='utf-8', errors='ignore') as f:
                    code = f.read()
                cli_node_map = extract_cli_node_from_code(code)
                for cli, nodes in cli_node_map.items():
                    all_cli_node_map[cli].update(nodes)
    return all_cli_node_map

def merge_cli_to_vtysh_json(cli_node_map, vtysh_json_path, output_path):
    """将CLI命令映射合并到vtysh.json中"""
    # 读取现有的vtysh.json
    with open(vtysh_json_path, 'r', encoding='utf-8') as f:
        vtysh_config = json.load(f)
    
    # 为每个节点创建命令列表
    node_commands = defaultdict(set)
    
    # 遍历CLI命令映射，按节点分组
    for cli, nodes in cli_node_map.items():
        # 提取CLI命令的第一个词作为基础命令
        base_command = cli.split()[0] if cli.strip() else "unknown"
        
        for node in nodes:
            if node in vtysh_config["nodes"]:
                node_commands[node].add(base_command)
    
    # 将命令添加到每个节点
    for node_name, node_info in vtysh_config["nodes"].items():
        if node_name in node_commands:
            # 转换为排序的列表
            commands_list = sorted(list(node_commands[node_name]))
            node_info["commands"] = commands_list
        else:
            # 如果没有找到命令，添加空列表
            node_info["commands"] = []
    
    # 添加CLI命令详细映射到配置中
    vtysh_config["cli_commands"] = {}
    for cli, nodes in cli_node_map.items():
        vtysh_config["cli_commands"][cli] = {
            "nodes": sorted(list(nodes)),
            "description": f"CLI命令: {cli}"
        }
    
    # 写入新的合并文件
    with open(output_path, 'w', encoding='utf-8') as f:
        json.dump(vtysh_config, f, ensure_ascii=False, indent=2)
    
    print(f"已合并配置输出到 {output_path}")
    return vtysh_config

if __name__ == "__main__":
    folder = "/home/ubuntu/sonic-buildimage/src/sonic-utilities/sonic-cli/vtysh"  # 修改为你的文件夹路径
    vtysh_json_path = "vtysh.json"
    output_path = "vtysh_merged.json"
    
    # 扫描代码生成CLI命令映射
    cli_node_map = scan_folder(folder)
    
    # 转换为可序列化的字典
    cli_output = []
    for cli, nodes in cli_node_map.items():
        cli_output.append({
            "cli": cli,
            "nodes": list(nodes)
        })
    
    # 写入CLI映射文件
    with open("cli_node_map.json", "w", encoding="utf-8") as f:
        json.dump(cli_output, f, ensure_ascii=False, indent=2)
    print("已输出到 cli_node_map.json")
    
    # 合并到vtysh.json
    if os.path.exists(vtysh_json_path):
        merge_cli_to_vtysh_json(cli_node_map, vtysh_json_path, output_path)
    else:
        print(f"警告: {vtysh_json_path} 文件不存在")
